"""This Module compare the two molecule's inner coords and make report
"""

from kuai.mol import *
from kuai.charges import dipole
from kuai.iotool import *
from kuai.inner import *
from kuai.kuaiext import *
from math import degrees, radians, sqrt, cos, pi
import sys

def atom_name(mol1, atom):
    return atom.symbol() + str(mol1.index(atom))

def diff_charges(mol1, mol2, output=sys.stdout):
    assert is_same_mol(mol1, mol2)
    if len(mol1.atom1) > 0:
        print_title("Compare of Charges")
        output.write("%6s %16s %10s %10s %10s %10s\n" \
                         %("", "Name", "E1 (e)", "E2 (e)", "Diff (e)", "Diff (%)"))
        max1 = 0.0
        msd1 = 0.0
        bias = 0.0
        for i in range(len(mol1.atoms)):
            a1 = mol1.atoms[i].partial_charge
            a2 = mol2.atoms[i].partial_charge
            abs_diff = a1-a2
            msd1 += abs_diff*abs_diff
            if abs(abs_diff) > abs(max1):
                max1 = abs_diff
            bias += abs_diff
            output.write("%6d %16s %10.6f %10.6f %10.6f\n" \
                         %(i+1, atom_name(mol1, mol1.atoms[i]), a1, a2, abs_diff))

        print_vertical_bar(level=1)
        output.write("%6s %16s %10s %10s %10.4f\n" %("", "Max", "", "", max1))
        if len(mol1.atoms) > 1:
            msd1 = sqrt(msd1/(len(mol1.atoms)-1))
            msd2 = sqrt(msd2/(len(mol1.atoms)-1))
            output.write("%6s %16s %10s %10s %10.4f\n" %("", "MSD", "", "", msd1))
        bias /= len(mol1.bonds)
        output.write("%6s %16s %10s %10s %10.4f\n" %("", "BIAS", "", "", bias))
        output.write("\n")

def diff_dipole(mol1, mol2, output=sys.stdout):
    if isinstance(mol1, Molecule):
        d1 = dipole(mo1)
    else:
        d1 = mol1
    d1 = dipole(mo1)
    if isinstance(mol2, Molecule):
        d2 = dipole(mo2)
    else:
        d2 = mol2
    print_title("Compare of Dipoles")
    output.write("%6s %16s %10s %10s %10s %10s\n" \
                         %("", "Name", "E1 (e*A)", "E2 (e*A)", "Diff (e*A)"))

    r = d1 - d2

    output.write("%6s %16s %10.4f %10.4f %10.4f\n" \
                         %("", "X", d1.x, d2.x, r.x))
    output.write("%6s %16s %10.4f %10.4f %10.4f\n" \
                         %("", "Y", d1.y, d2.y, r.y))
    output.write("%6s %16s %10.4f %10.4f %10.4f\n" \
                         %("", "Z", d1.z, d2.z, r.z))
    print_vertical_bar(level=1)
    output.write("%6s %16s %10s %10s %10.4f\n" %("", "R", "", "", abs(r)))
    output.write("\n")


def bond_name(mol, bond):
    return bond.atom1.symbol() + str(mol.index(bond.atom1)) + "-" \
            + bond.atom2.symbol() + str(mol.index(bond.atom2))

def diff_bonds(mol1, mol2, output=sys.stdout):
    assert is_same_mol(mol1, mol2)
    if len(mol1.bonds) > 0:
        print_title("Compare of Bonds")
        output.write("%6s %16s %10s %10s %10s %10s\n" \
                         %("", "Name", "R1 (A)", "R2 (A)", "Diff (A)", "Diff (%)"))
        max1 = 0.0
        max2 = 0.0
        msd1 = 0.0
        msd2 = 0.0
        bias = 0.0
        for i in range(len(mol1.bonds)):
            b1 = mol1.bonds[i]
            b2 = mol2.bonds[i]
            r1 = abs(b1.atom1.coords - b1.atom2.coords)
            r2 = abs(b2.atom1.coords - b2.atom2.coords)
            abs_diff = r1-r2
            rel_diff = (r1-r2)/(r1+r2)*200.0
            msd1 += abs_diff*abs_diff
            msd2 += rel_diff*rel_diff
            if abs(abs_diff) > abs(max1):
                max1 = abs_diff
            if abs(rel_diff) > abs(max2):
                max2 = rel_diff
            bias += abs_diff
            output.write("%6d %16s %10.4f %10.4f %10.4f %10.4f\n" \
                         %(i+1, bond_name(mol1, b1), r1, r2, abs_diff, rel_diff))

        print_vertical_bar(level=1)
        output.write("%6s %16s %10s %10s %10.4f %10.4f\n" %("", "Max", "", "", max1, max2))
        if len(mol1.bonds) > 1:
            msd1 = sqrt(msd1/(len(mol1.bonds)-1))
            msd2 = sqrt(msd2/(len(mol1.bonds)-1))
            output.write("%6s %16s %10s %10s %10.4f %10.4f\n" %("", "MSD", "", "", msd1, msd2))
        bias /= len(mol1.bonds)
        output.write("%6s %16s %10s %10s %10.4f %10s\n" %("", "BIAS", "", "", bias, ""))
        output.write("\n")

def angle_name(mol, angle):
    return angle[0].symbol() + str(mol.index(angle[0])) + "-" \
            + angle[1].symbol() + str(mol.index(angle[1])) + "-" \
            + angle[2].symbol() + str(mol.index(angle[2]))


def diff_angles(mol1, mol2, output=sys.stdout):
    assert is_same_mol(mol1, mol2)
    angle1 = list_angle_atoms(mol1)
    angle2 = list_angle_atoms(mol2)
    assert len(angle1) == len(angle2)
    if len(angle1) > 0:
        print_title("Compare of Angles")
        output.write("%6s %16s %10s %10s %10s %10s\n" \
                         %("", "Name", "A1 (D)", "A2 (D)", "Diff (D)", "Diff (%)"))
        max1 = 0.0
        max2 = 0.0
        msd1 = 0.0
        msd2 = 0.0
        bias = 0.0
        for i in range(len(angle1)):
            a1 = angle(angle1[i][0].coords, angle1[i][1].coords, angle1[i][2].coords)
            a2 = angle(angle2[i][0].coords, angle2[i][1].coords, angle2[i][2].coords)
            a1 = degrees(a1)
            a2 = degrees(a2)
            abs_diff = a1-a2
            rel_diff = (a1-a2)/(a1+a2)*200.0
            msd1 += abs_diff*abs_diff
            msd2 += rel_diff*rel_diff
            if abs(abs_diff) > abs(max1):
                max1 = abs_diff
            if abs(rel_diff) > abs(max2):
                max2 = rel_diff
            bias += abs_diff
            output.write("%6d %16s %10.4f %10.4f %10.4f %10.4f\n" \
                         %(i+1, angle_name(mol1, angle1[i]), a1, a2, abs_diff, rel_diff))
        print_vertical_bar(level=1)
        output.write("%6s %16s %10s %10s %10.4f %10.4f\n" %("", "Max", "", "", max1, max2))
        if len(angle1) > 1:
            msd1 = sqrt(msd1/(len(angle1)-1))
            msd2 = sqrt(msd2/(len(angle1)-1))
            output.write("%6s %16s %10s %10s %10.4f %10.4f\n" %("", "MSD", "", "", msd1, msd2))
        bias /= len(mol1.bonds)
        output.write("%6s %16s %10s %10s %10.4f %10s\n" %("", "BIAS", "", "", bias, ""))
        output.write("\n")

def torsion_name(mol, torsion):
    return torsion[0].symbol() + str(mol.index(torsion[0])) + "-" \
            + torsion[1].symbol() + str(mol.index(torsion[1])) + "-" \
            + torsion[2].symbol() + str(mol.index(torsion[2])) + "-" \
            + torsion[3].symbol() + str(mol.index(torsion[3]))

TWO_DEGREES = radians(2)
def diff_torsions(mol1, mol2, output=sys.stdout):
    assert is_same_mol(mol1, mol2)
    angle1 = list_torsion_atoms(mol1)
    angle2 = list_torsion_atoms(mol2)
    assert len(angle1) == len(angle2)
    if len(angle1) > 0:
        print_title("Compare of Torsion")
        output.write("%6s %16s %10s %10s %10s\n" \
                         %("", "Name", "T1 (D)", "T2 (D)", "Diff (D)"))
        max1 = 0.0
        msd1 = 0.0
        bias = 0.0
        for i in range(len(angle1)):
            a1 = angle(angle1[i][0].coords, angle1[i][1].coords, angle1[i][2].coords)
            a2 = angle(angle1[i][1].coords, angle1[i][2].coords, angle1[i][3].coords)
            if abs(a1-pi) < TWO_DEGREES or abs(a2-pi) < TWO_DEGREES:
                continue

            a1 = angle(angle2[i][0].coords, angle2[i][1].coords, angle2[i][2].coords)
            a2 = angle(angle2[i][1].coords, angle2[i][2].coords, angle2[i][3].coords)
            if abs(a1-pi) < TWO_DEGREES or abs(a2-pi) < TWO_DEGREES:
                continue

            a1 = torsion(angle1[i][0].coords, angle1[i][1].coords, angle1[i][2].coords, angle1[i][3].coords)
            a2 = torsion(angle2[i][0].coords, angle2[i][1].coords, angle2[i][2].coords, angle2[i][3].coords)
            a1 = degrees(a1)
            a2 = degrees(a2)
            abs_diff = a1-a2
            if abs_diff <= -180:
                abs_diff += 360
            elif abs_diff >= 180:
                abs_diff -= 360
            msd1 += abs_diff*abs_diff
            if abs(abs_diff) > abs(max1):
                max1 = abs_diff
            bias += abs_diff
            output.write("%6d %16s %10.4f %10.4f %10.4f\n" \
                         %(i+1, torsion_name(mol1, angle1[i]), a1, a2, abs_diff))
        print_vertical_bar(level=1)
        output.write("%6s %16s %10s %10s %10.4f\n" %("", "Max", "", "", max1))
        if len(angle1) > 1:
            msd1 = sqrt(msd1/(len(angle1)-1))
            output.write("%6s %16s %10s %10s %10.4f\n" %("", "MSD", "", "", msd1))
        bias /= len(mol1.bonds)
        output.write("%6s %16s %10s %10s %10.4f\n" %("", "BIAS", "", "", bias))
        output.write("\n")

def diff_cos_torsions(mol1, mol2, output=sys.stdout):
    assert is_same_mol(mol1, mol2)
    angle1 = list_torsion_atoms(mol1)
    angle2 = list_torsion_atoms(mol2)
    assert len(angle1) == len(angle2)
    if len(angle1) > 0:
        print_title("Compare COS of Torsion")
        output.write("%6s %16s %10s %10s %10s %10s\n" \
                         %("", "Name", "T1   ", "T2   ", "Diff   ", ""))
        max1 = 0.0
        msd1 = 0.0
        bias = 0.0
        for i in range(len(angle1)):
            a1 = torsion(angle1[i][0].coords, angle1[i][1].coords, angle1[i][2].coords, angle1[i][3].coords)
            a2 = torsion(angle2[i][0].coords, angle2[i][1].coords, angle2[i][2].coords, angle1[i][3].coords)
            a1 = cos(a1)
            a2 = cos(a2)
            abs_diff = a1-a2
            msd1 += abs_diff*abs_diff
            if abs(abs_diff) > abs(max1):
                max1 = abs_diff
            bias += abs_diff
            output.write("%6d %16s %10.4f %10.4f %10.4f\n" \
                         %(i+1, torsion_name(mol1, angle1[i]), a1, a2, abs_diff))
        print_vertical_bar(level=1)
        output.write("%6s %16s %10s %10s %10.4f\n" %("", "Max", "", "", max1))
        if len(angle1) > 1:
            msd1 = sqrt(msd1/(len(angle1)-1))
            output.write("%6s %16s %10s %10s %10.4f\n" %("", "MSD", "", "", msd1))
        bias /= len(mol1.bonds)
        output.write("%6s %16s %10s %10s %10.4f %10s\n" %("", "BIAS", "", "", bias, ""))
        output.write("\n")

def compare_mol(mol1, mol2, output=sys.stdout):
    print_block("Compare of "+ mol1.name)
    output.write("\n")
    diff_bonds(mol1, mol2, output)
    diff_angles(mol1, mol2, output)
    diff_torsions(mol1, mol2, output)
    output.write("\n\n\n\n")

def digest_compare(input, output):
    def count(file):
        n = 0
        for i in file:
            if i.startswith("    ----"):
                break
            else:
                n += 1
        n -= 2
        result = [n]
        tokens = file.next().split()[1:]
        result += tokens
        tokens = file.next().split()[1:]
        result += tokens
        tokens = file.next().split()[1:]
        result += tokens
        return result

    def save(output, name, bonds, angles, torsions):
        if name:
            output.write(name)
        output.write(",")
        if bonds:
            output.write(",".join([str(i) for i in bonds])+",")
        else:
            output.write(","*6)
        if angles:
            output.write(",".join([str(i) for i in angles])+",")
        else:
            output.write(","*6)
        if torsions:
            output.write(",".join([str(i) for i in torsions])+",")
        else:
            output.write(","*4)
        output.write("\n")

    def parse_name(file):
        file.next()
        i = file.next()
        i1 = i.index("Compare of ")+len("Compare of ")
        i2 = i.index("*", i1)
        name = i[i1:i2].strip()
        file.next()
        file.next()
        return name

    output.write("name,nbond,max(A),max(R),msd(A),msd(R),Avg,nAngles,max(A),max(R),msd(A),msd(R),Avg,nTorsion,max(A),msd(A),Avg,\n")
    name = None
    for i in input:
        if i.startswith("****"):
            if name != None:
                save(output, name, bonds, angles, torsions)
            name = parse_name(input)
            bonds = None
            angles = None
            torsions = None
        elif i.find("Compare of Bonds") != -1:
            bonds = count(input)
        elif i.find("Compare of Angles") != -1:
            angles = count(input)
        elif i.find("Compare of Torsion") != -1:
            torsions = count(input)
        else:
            pass

def filter_compare(input, output, title):
    start = False
    for i in input:
        if start:
            if i.find('----------------------------------------------------------------') != -1:
                start = False
            else:
                output.write(i)
        elif i.find(title) != -1:
            start = True
            input.next()
            input.next()
        #The following three lines were added to output molecule the terms belong to
        #by Chengtao 2009-6-17
        elif len(i.split()) > 2:
            if i.split()[0] == '*' and i.split()[1] == 'Compare':
                file_name = i.split()[3]


def compare_hessian(hessian1, hessian2, output=sys.stdout):
    assert len(hessian1) == len(hessian2)
    print_title("Compare of Hessian")
    output.write("%6s %16s %10s %10s %10s %10s\n" \
                     %("", "Name", "E1 (e)", "E2 (e)", "Diff (e)", "Diff (%)"))
    max1 = 0.0
    msd1 = 0.0
    bias = 0.0
    for i in range(len(hessian1)):
        a1 = hessian1[i]
        a2 = hessian2[i]
        abs_diff = a1-a2
        msd1 += abs_diff*abs_diff
        if abs(abs_diff) > abs(max1):
            max1 = abs_diff
        bias += abs_diff
        output.write("%6d %16s %10.3f %10.3f %10.3f\n" \
                     %(i+1, "", a1, a2, abs_diff))

    print_vertical_bar(level=1)
    output.write("%6s %16s %10s %10s %10.3f\n" %("", "Max", "", "", max1))
    assert len(hessian1) > 1
    msd1 = sqrt(msd1/(len(hessian1)-1))
    output.write("%6s %16s %10s %10s %10.3f\n" %("", "MSD", "", "", msd1))
    bias /= len(hessian1)
    output.write("%6s %16s %10s %10s %10.3f\n" %("", "BIAS", "", "", bias))
    output.write("\n")
